from seal_recognition.seal_reg import seal_reg
import requests
import json
import fitz
import os
import docx
from docx.shared import Pt
from docx.oxml.ns import qn
from spire.doc import *
from spire.doc.common import *


directory_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
basepath = os.path.join(directory_path, 'file')
graphpath = os.path.join(directory_path, 'graph')
sealpath = os.path.join(directory_path, 'seal_output')
outputpath = os.path.join(directory_path, 'final_output')

def pdf_to_text(basepath):
    PDFText = {}
    for file in os.listdir(basepath):
        file_path = os.path.join(basepath, file)
        # 检查是否为pdf文件
        if os.path.isfile(file_path) and file_path.endswith('.pdf'):
            filename = file.split('.')[0]
            allText = []
            doc = fitz.open(file_path)

            # 读取PDF文件内容
            for page in doc:
                text = page.get_text("text")
                # print(text)
                textSplit = text.split('\n')
                # print(textSplit)
                allText.extend(textSplit)

            textString = ''.join(allText)
            # print(textString)
            # print(len(textString))

            PDFText[filename] = textString
    return PDFText

def docx_to_text(basepath):
    DOCXText = {}
    for file in os.listdir(basepath):
        file_path = os.path.join(basepath, file)
        # 检查是否为docx文件
        if os.path.isfile(file_path) and file_path.endswith('.docx'):
            filename = file.split('.')[0]
            doc = docx.Document(file_path)
            allText = []
            for para in doc.paragraphs:
                allText.append(para.text)
            textString = '\n'.join(allText)
            DOCXText[filename] = textString
        return DOCXText

def pdf_to_image(basepath, graphpath):
    for file in os.listdir(basepath):
        file_path = os.path.join(basepath, file)
        # 检查是否为pdf文件
        if os.path.isfile(file_path) and file_path.endswith('.pdf'):
            filename = file.split('.')[0]
            doc = fitz.open(file_path)
            # 输出最后两页的图片
            new_graph_path = os.path.join(graphpath, filename)
            try:
                # 创建新文件夹
                os.mkdir(new_graph_path)
                print(f"成功在 {graphpath} 下创建文件夹 {filename}")
            except FileExistsError:
                print(f"文件夹 {graphpath} 已存在。")
            except PermissionError:
                print("没有权限创建文件夹。")
            except Exception as e:
                print(f"创建文件夹时出现错误: {e}")

            page_last = doc[-1]
            pix1 = page_last.get_pixmap()
            pix1.save(os.path.join(new_graph_path, 'last_page.png'))
            if len(doc)>1:
                page_second_last = doc[-2]
                pix2 = page_second_last.get_pixmap()
                pix2.save(os.path.join(new_graph_path, 'second_page.png'))

def markdown_to_pdf(input_md_path, output_pdf_path):
    doc = Document()
    # 加载一个 Markdown 文件
    doc.LoadFromFile(input_md_path, FileFormat.Markdown)
    # 将文件保存为 PDF 文档
    doc.SaveToFile(output_pdf_path, FileFormat.PDF)
    doc.Dispose()


def save_to_markdown(text1, seals_num):
    # 分割文字，找到关键标题位置
    parts = text1.split("数据资产合规风险点：")
    basic_info = parts[0]
    compliance_info = "数据资产合规风险点：" + parts[1]
    # 添加 Markdown 最大标题字号
    basic_info_title = "# " + basic_info.split("\n", 1)[0]
    basic_info_content = basic_info.split("\n", 1)[1]
    compliance_info_title = "# " + compliance_info.split("\n", 1)[0]
    compliance_info_content = compliance_info.split("\n", 1)[1]
    # 组合 Markdown 内容
    markdown_text = basic_info_title + "\n" + basic_info_content + "\n" + compliance_info_title + "\n" + compliance_info_content
    # 展示印章的数量和路径
    seal_title = "# " + "数据资产格式风险点：\n"
    seal_content = "## 盖章情况: \n"
    for key, value in seals_num.items():
        seal_content += f"### 企业上传的{key}中有{value}个印章\n"
        for _, temp, image_list in os.walk(os.path.join(sealpath, key)):
            for image_name in image_list:
                if image_name.split('_')[0]=='last' and image_name.split('_')[-1]=='layout.jpg':
                    # 定义图片路径和描述
                    image_path = os.path.join('../seal_output', key, image_name)
                    image_description = f'{image_name}'
                    # 构建 Markdown 格式的图片插入语句
                    markdown_image = f"![{image_description}]({image_path})" + "\n"
                    seal_content += markdown_image
    markdown_text += '\n' + seal_title + seal_content
    # 保存为 Markdown 文件
    md_name = "test4.md"
    with open(os.path.join(outputpath, md_name), "w", encoding="utf-8") as file:
        file.write(markdown_text)
    print(f"成功保存{md_name}")
    # pdf_name = "test3.pdf"
    # markdown_to_pdf(os.path.join(outputpath, md_name), os.path.join(outputpath, pdf_name))


def newFile(new_file_path, result):
    # 先尝试打开已有文档，如果不存在则创建新文档
    if os.path.exists(new_file_path):
        doc = docx.Document(new_file_path)
    else:
        doc = docx.Document()

    # 遍历字典，将每个键值对作为标题和段落内容添加到文档中
    for title, content in result.items():
        heading_1 = doc.add_heading(title, level=1)
        set_song_font(heading_1)

        if isinstance(content, str):
            paragraph = doc.add_paragraph(content)
            set_song_font(paragraph)
        else:
            for title2, content2 in content.items():
                heading_2 = doc.add_heading(title2, level=2)
                set_song_font(heading_2)

                if isinstance(content2, str):
                    paragraph = doc.add_paragraph(content2)
                    set_song_font(paragraph)
                else:
                    for title3, content3 in content2.items():
                        heading_3 = doc.add_heading(title3, level=3)
                        set_song_font(heading_3)

                        paragraph = doc.add_paragraph(content3)
                        set_song_font(paragraph)

    # 保存文档
    doc.save(new_file_path)
    print(f"成功保存至{new_file_path}")


def set_song_font(element, style = "paragraph"):
    """
    设置元素（段落或标题）的字体为宋体
    :param element: 段落或标题对象
    """
    if style == "heading_1":
        for run in element.runs:
            run.font.name = '等线'
            run.font.size = Pt(14)
            run._element.rPr.rFonts.set(qn('w:eastAsia'), '等线')
    elif style == "heading_2":
        for run in element.runs:
            run.font.name = '等线'
            run.font.size = Pt(12)
            run._element.rPr.rFonts.set(qn('w:eastAsia'), '等线')
    elif style == "heading_3":
        for run in element.runs:
            run.font.name = '等线'
            run.font.size = Pt(10.5)
            run._element.rPr.rFonts.set(qn('w:eastAsia'), '等线')
    else:
        for run in element.runs:
            run.font.name = '等线'
            run.font.size = Pt(10.5)
            run._element.rPr.rFonts.set(qn('w:eastAsia'), '等线')

def extract_hegui_sections(basepath = basepath, file_path = None):
    if not file_path:
        for file in os.listdir(basepath):
            file_path = os.path.join(basepath, file)
            # 检查是否为pdf文件
            if os.path.isfile(file_path) and file_path.endswith('.docx'):
                break
    doc = docx.Document(file_path)
    target_sections = [
        "数据产权证书合规审查及建议",
        "数据资源入表资源范围合规及建议",
        "数据治理合规性审查及建议"
    ]
    result = {}
    current_section = None
    stop_section = "、合规改进措施"  # 定义停止收集内容的标题

    for para in doc.paragraphs:
        para_text = para.text.strip()
        # 检查是否为目标标题
        if para_text.split('、')[-1] in target_sections:
            current_section = para_text
            result[current_section] = []
        # 检查是否遇到停止收集内容的标题
        elif para_text[1:] == stop_section:
            current_section = None
        # 如果当前有正在收集内容的目标标题，且段落文本不为空，则添加到结果中
        elif current_section and para_text:
            result[current_section].append(para_text)

    for section, content in result.items():
        result[section] = '\n'.join(content)
    result_str = '\n\n'.join([f"{k}\n{v}" for k, v in result.items()])
    return result_str

def right_corresponding(rights):
    corr_rights = []
    for right in rights:
        if right == '数据持有权':
            pass
        elif right == '数据使用权': 
            corr_rights.append('自我使用')
            corr_rights.append('可经营')
        elif right == '数据经营权':
            corr_rights.append('可转让')
            corr_rights.append('可收益')
            corr_rights.append('可经营')
    corr_rights = list(set(corr_rights))
    return corr_rights

if __name__ == '__main__':
    # 入表智能体
    api_key = "app-DZYVdGjAsHBYKeqbQwD0hSDA"
    # 改进点：可以根据合同文本长度来判断是否需要从经过一个大模型变为经过多个大模型
    api_url = "http://10.1.31.201:8080/v1/workflows/run"

    # Request headers
    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': 'application/json',
    }
    basic_info = {"resource":"自行生产", "company_name":"上海云素网络科技有限公司", "right_type":["数据持有权", "数据使用权", "数据经营权"], "right_content":["自我使用", "可转让", "可收益", "可经营"]}
    corr_rights = right_corresponding(basic_info["right_type"])

    # 上传的产权证明材料
    appendixes = pdf_to_text(basepath)
    input1 = f"""
    你是一个数据资产登记的审核员，你需要根据用户（企业）提供的数据资产信息，从下面的角度来分析为该企业进行数据资产登记有哪些风险点。
    （1）在企业自行填写的基本信息中，数据来源为“{basic_info["resource"]}”，根据其提供的产权证明材料，判断企业填写的数据来源是否正确，是否存在风险；
    （2）在企业自行填写的基本信息中，数据资产的产权类型有{basic_info["right_type"]}，这几种权益的概念分别是：
    -数据持有权，是指权利人自行持有或委托他人代为持有合法获取的数据的权利，旨在防范他人非法违规窃取、篡改、泄露或者破坏持有权人持有的数据。
    -数据使用权，是指权利人通过加工、聚合、分析等方式，将数据用于优化生产经营、形成衍生数据等的权利。一般来说，使用权是权利人在不对外提供数据的前提下，将数据用于内部使用的权利。
    -数据经营权，是指权利人通过转让、许可、出资或者设立担保等有偿或无偿的方式对外提供数据的权利。
    结合上述三种产权的概念，根据企业提供的产权证明材料判断该企业是否分别拥有这三种权益，若没有，则给出对应的潜在风险；
    （3）在企业自行填写的基本信息中，数据资产的产权具体内容有{basic_info["right_content"]}，这几种权益的概念分别是：
    -自我使用：如果有权直接使用其数据，以满足个人或组织内部的需求，则代表拥有数据加工使用权。
    -可经营：如果有权依法合规进行数据汇聚、治理、加工处理，提供数据资源、产品与服务，则代表拥有数据经营权。
    -可转让：如果有权将数据的产权完全转移给第三方，包括出售、赠与等，则代表拥有数据产品经营权。
    -可收益：如果数据能够被用于商业活动，并且有潜力创造经济价值，则代表拥有数据产品经营权。
    而根据专家判断出的结果为该企业具有的产权具体内容为：{corr_rights}，结合相关定义，针对缺少的产权内容，给出企业的潜在风险。
    （4）企业基本信息中填写的组织机构名称为“{basic_info["company_name"]}”，核查该名称是否与下述的产权证明材料及数据资产登记承诺书中的名称一致，若不一致，给出对应的法律风险；
    
    该企业的产权证明材料（数据转让合同或数据许可授权合同等合同）和数据资产登记承诺书如下：
    {appendixes}
    """
    # print(input1)
    # print(len(input1))
    input2 = "根据该企业的数据资产合规法律意见书，总结其中提到的潜在风险\n法律意见书：\n" + extract_hegui_sections()
    # print(input2)

    # 储存工作流输出
    workflow_output = ""
    payload = {
                    "inputs": {"input1": input1, "input2": input2},
                    "response_mode": "blocking",
                    "user": "a918@buaa"
                }
    response = requests.post(api_url, headers=headers, json=payload)
    # print(response.text)
    # 检查响应状态码
    if response.status_code == 200 :
        # 提取 Dify 响应中的消息内容
        dify_response = response.json()
        # 假设 Dify 的响应中有 answer 字段作为回复内容，根据实际情况修改
        # 聊天
        # assistant_message = dify_response.get('answer')
        # 工作流
        assistant_message = dify_response.get('data').get('outputs').get('result')
        print(f"assistant_message: {assistant_message}")
        workflow_output = assistant_message

    # 判断盖章数量，并生成识别图片
    pdf_to_image(basepath,graphpath)
    seals_num = seal_reg(graphpath, sealpath)
    save_to_markdown(workflow_output, seals_num)
    
    # Send the POST request with streaming
    # with requests.post(api_url, headers=headers, json=payload, stream=True) as response:
    #     # Check the response status code
    #     if response.status_code == 200:
    #         print("Request successful")
    #         # Iterate over the lines in the response
    #         for line in response.iter_lines():
    #             if line:
    #                 try:
    #                     # Decode the line if it's in bytes
    #                     if isinstance(line, bytes):
    #                         line = line.decode('utf-8')
    #                     if line.startswith("event: ping"):
    #                         continue
    #                     # Remove the "data: " prefix
    #                     if line.startswith("data: "):
    #                         line = line[len("data: "):]
    #                     # Parse the JSON string
    #                     json_data = json.loads(line)
    #                     try:
    #                         text = json_data.get("data").get("text")
    #                         if text:
    #                             print(text, end='', flush=True)
    #                     except:
    #                         pass
    #                 except Exception as e:
    #                     print(f"Error decoding line: {line}")
    #     else:
    #         print(f"Request failed with status code: {response.status_code}")
    #         print("Error:", response.text)